package com.tiny.work.test.instrument;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtBehavior;
import javassist.CtClass;
import javassist.expr.ExprEditor;
import javassist.expr.MethodCall;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

public class PerfMonXformer implements ClassFileTransformer{
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        byte[] transformed = null;

        System.out.println(className + " is transforming");

        ClassPool pool = ClassPool.getDefault();
        CtClass ctClass = null;

        try {
            ctClass = pool.makeClass(new ByteArrayInputStream(classfileBuffer));

            if (ctClass.isInterface() == false) {
                CtBehavior[] method = ctClass.getDeclaredBehaviors();
                for (int i =0; i< method.length; i++) {
                    doMethod(method[i]);
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        } catch (CannotCompileException e) {
            e.printStackTrace();

        } finally {
            if (ctClass != null) {
                ctClass.detach();
            }
        }
        return transformed;
    }
    private void doMethod(CtBehavior behavior) throws CannotCompileException {
        behavior.instrument(new ExprEditor(){
            @Override
            public void edit(MethodCall m) throws CannotCompileException {
                super.edit(m);
                m.replace("{ long stime = System.nanoTime(); $_ = $proceed($$); System.out.println(\""
                                + m.getClassName()+"."+m.getMethodName()
                                + ":\"+(System.nanoTime()-stime));}");
            }
        });
    }
 }
